// scene.ts 
import { ISize } from "./interface/commTypes";
import { IGeometry, IGeometryBuffer } from "./interface/geometry";
import PipelineObject from "./interface/pipelineObject";
import PipelineLightObject from "./interface/pipelineLightObject";
import { geometryToGPUBuffer } from "./util/geometryToGPUBuffer";
import { createGBuffer } from "./pipeline/pipelineGeometry";
import GBufferTextures from "./interface/gBufferTextures";
import Camera from "./camera";
import ShadowPipelineInfo from './interface/shadowPipelineInfo'
import { createShadowPipeline } from "./pipeline/pipelineShadow";

export default class Scene {

    private device: GPUDevice;
    private context: GPUCanvasContext;
    private format: GPUTextureFormat;
    private size: ISize;
    private lightingVertexBuffer: GPUBuffer;
    private lightPipeline: PipelineLightObject | undefined;
    private gbuffer: GBufferTextures | undefined;
    camera: Camera;
    shadowPipeline: ShadowPipelineInfo | undefined;
    private _enableShadow: boolean = false;


    // Getter for enableShadow
    get enableShadow(): boolean {
        return this._enableShadow;
    }

    // Setter for enableShadow
    set enableShadow(value: boolean) {

        this._enableShadow = value;
        // 在这里可以添加其他逻辑，例如创建或销毁阴影管道
        const { lightPipeline } = this;
        if (lightPipeline) {

            // 更新缓冲区
            const shadowEnabled = new Uint32Array([0]); // 1 表示 true，0 表示 false
            if (this._enableShadow) {
                shadowEnabled[0] = 1;
            }

            this.device.queue.writeBuffer(lightPipeline.shadowOptionsBuffer, 0, shadowEnabled);
        }
    }

    constructor(
        device: GPUDevice,
        context: GPUCanvasContext,
        format: GPUTextureFormat,
        size: ISize
    ) {
        this.device = device;
        this.context = context;
        this.format = format;
        this.size = size;
        this.enableShadow = false;

        // 创建光照阶段的顶点数据（全屏四边形）
        const lightingVertices = new Float32Array([
            -1.0, -1.0, // 左下角
            1.0, -1.0, // 右下角
            1.0, 1.0, // 右上角
            1.0, 1.0, // 右上角
            -1.0, 1.0, // 左上角
            -1.0, -1.0, // 左下角
        ]);

        this.lightingVertexBuffer = device.createBuffer({
            size: lightingVertices.byteLength,
            usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
        });
        device.queue.writeBuffer(this.lightingVertexBuffer, 0, lightingVertices);

        this.camera = new Camera();
        createShadowPipeline(device).then((pipeline) => {
            this.shadowPipeline = pipeline;
        });
    }

    public setLightingPipeline(pipeline: PipelineLightObject) {
        this.lightPipeline = pipeline;

        this.enableShadow = this.enableShadow;
    }

    // 方法：渲染场景
    public render() {

        const { instancePipelineMap, lightPipeline, shadowPipeline, enableShadow } = this;
        if (!shadowPipeline) {
            return;
        }

        var pps = instancePipelineMap.keys().toArray();
        if (pps.length < 1 || !lightPipeline) {
            return;
        }

        const { device, context } = this;
        const commandEncoder = device.createCommandEncoder()

        let firstRender = true;

        if (enableShadow) {
            for (let i = 0, il = pps.length; i < il; i++) {
                const pipeline = pps[i];
                const geometries = instancePipelineMap.get(pipeline) ?? [];
                if (geometries && geometries.length > 0 && pipeline) {
                    let loadOp: GPULoadOp = firstRender ? 'clear' : 'load';
                    this._renderShadowMap(commandEncoder, loadOp, geometries, pipeline);
                    firstRender = false;
                }
            }
        }

        firstRender = true;
        for (let i = 0, il = pps.length; i < il; i++) {
            const pipeline = pps[i];
            const geometries = instancePipelineMap.get(pipeline) ?? [];
            if (geometries && geometries.length > 0 && pipeline) {
                let loadOp: GPULoadOp = firstRender ? 'clear' : 'load';
                this.renderPipeline(commandEncoder, loadOp, pipeline, geometries);
                firstRender = false;
            }
        }

        const gBufferTextures = pps[0].gBufferTextures;
        if (!gBufferTextures || !gBufferTextures.position || !gBufferTextures.normal || !gBufferTextures.color) {
            return;
        }

        // 创建绑定组
        const bindGroup = device.createBindGroup({
            label: 'Uniform Group with gbuffer and shadowmap',
            layout: lightPipeline.pipeline.getBindGroupLayout(0), // 使用管线的绑定组布局
            entries: [
                {
                    binding: 0,
                    resource: gBufferTextures.position.createView(),
                },
                {
                    binding: 1,
                    resource: gBufferTextures.normal.createView(),
                },
                {
                    binding: 2,
                    resource: gBufferTextures.color.createView(),
                },
                {
                    binding: 3,
                    resource: {
                        buffer: shadowPipeline.lightProjectionBuffer
                    }
                },
                {
                    binding: 4,
                    resource: shadowPipeline.depthView,
                },
                {
                    binding: 5,
                    resource: device.createSampler({
                        compare: 'less',
                    })
                },
                {
                    binding: 6,
                    resource: {
                        buffer: lightPipeline.shadowOptionsBuffer
                    }
                },
            ],
        });

        const lightingPass = commandEncoder.beginRenderPass({
            colorAttachments: [
                {
                    view: context.getCurrentTexture().createView(),
                    loadOp: 'clear',
                    clearValue: { r: 0.0, g: 0.0, b: 0.0, a: 1.0 },
                    storeOp: 'store',
                },
            ],
        });
        lightingPass.setPipeline(lightPipeline.pipeline);
        lightingPass.setBindGroup(0, bindGroup); // 绑定 G-Buffer 纹理
        lightingPass.setBindGroup(1, lightPipeline.lightGroup);
        lightingPass.setVertexBuffer(0, this.lightingVertexBuffer);
        lightingPass.draw(6);
        lightingPass.end();

        // webgpu run in a separate process, all the commands will be executed after submit
        device.queue.submit([commandEncoder.finish()])
    }

    private instanceMap: Map<IGeometry, number> = new Map();
    private instancePipelineMap: Map<PipelineObject, Array<IGeometry>> = new Map();
    private instanceBufferMap: Map<IGeometry, IGeometryBuffer> = new Map();
    public addInstance(geometry: IGeometry, num: number) {
        if (!geometry.pipeline) {
            return;
        }

        const { instanceBufferMap, instanceMap, instancePipelineMap } = this;
        if (instanceMap.has(geometry)) {
            const onum = instanceMap.get(geometry);
            if (onum === undefined) {
                instanceMap.set(geometry, num);
            } else {
                instanceMap.set(geometry, onum + num);
            }
        } else {
            instanceMap.set(geometry, num);
        }

        if (!instanceBufferMap.has(geometry)) {
            instanceBufferMap.set(geometry, geometryToGPUBuffer(this.device, geometry));
        }

        if (!instancePipelineMap.has(geometry.pipeline)) {
            instancePipelineMap.set(geometry.pipeline, [])
        }

        instancePipelineMap.get(geometry.pipeline)?.push(geometry)

        return this;
    }

    public removeInstance(geometry: IGeometry) {
        const { instanceMap, instancePipelineMap, instanceBufferMap } = this;
        if (instanceMap.has(geometry)) {
            instanceMap.delete(geometry);
            const geometries = instancePipelineMap.get(geometry.pipeline);
            geometries?.splice(geometries?.indexOf(geometry), 1);
            if (geometries?.length === 0) {
                instancePipelineMap.delete(geometry.pipeline);
                instanceBufferMap.delete(geometry);
            }
        }
        return this;
    }

    private _renderShadowMap(
        commandEncoder: GPUCommandEncoder,
        loadOp: GPULoadOp,
        geometries: Array<IGeometry>,
        geomPipeline: PipelineObject) {
        const pipelineObj = this.shadowPipeline;
        if (!pipelineObj) {
            return;
        }

        const shadowPassDescriptor: GPURenderPassDescriptor = {
            colorAttachments: [],
            depthStencilAttachment: {
                view: pipelineObj.depthView,
                depthClearValue: 1.0,
                depthLoadOp: loadOp,
                depthStoreOp: 'store',
            }
        }


        const passEncoder = commandEncoder.beginRenderPass(shadowPassDescriptor)
        passEncoder.setPipeline(pipelineObj.pipeline)

        if (!geomPipeline.shadowGroup) {
            const { device } = this;
            geomPipeline.shadowGroup = device.createBindGroup({

                label: 'Uniform Group with matrix',
                layout: pipelineObj.pipeline.getBindGroupLayout(0),
                entries: [
                    {
                        binding: 0,
                        resource: {
                            buffer: geomPipeline.modelViewBuffer
                        }
                    },
                    {
                        binding: 1,
                        resource: {
                            buffer: pipelineObj.lightProjectionBuffer
                        }
                    },
                    {
                        binding: 2,
                        resource: {
                            buffer: geomPipeline.colorBuffer
                        }
                    },
                ]
            });
        }

        passEncoder.setBindGroup(0, geomPipeline.shadowGroup)

        const { instanceBufferMap, instanceMap } = this;

        let renderedNum = 0;
        for (let i = 0, il = geometries.length; i < il; i++) {
            const geometry = geometries[i];
            const buffer = instanceBufferMap.get(geometry);
            if (buffer) {
                const num = instanceMap.get(geometry) ?? 0;
                if (num > 0) {
                    passEncoder.setVertexBuffer(0, buffer.vertex)
                    passEncoder.setIndexBuffer(buffer.index, 'uint16')
                    passEncoder.drawIndexed(geometry.indexCount, num, 0, 0, renderedNum)
                    renderedNum += num;
                }
            }
        }

        passEncoder.end()
    }

    private renderPipeline(
        commandEncoder: GPUCommandEncoder,
        loadOp: GPULoadOp,
        pipeline: PipelineObject,
        geometries: Array<IGeometry>) {
        const { gBufferTextures } = pipeline;
        const renderPassDescriptor: GPURenderPassDescriptor = {
            colorAttachments: [
                {
                    view: gBufferTextures.position.createView(),
                    loadOp: loadOp,
                    clearValue: { r: 0.0, g: 0.0, b: 0.0, a: 1.0 },
                    storeOp: 'store',
                },
                {
                    view: gBufferTextures.normal.createView(),
                    loadOp: loadOp,
                    clearValue: { r: 0.0, g: 0.0, b: 0.0, a: 1.0 },
                    storeOp: 'store',
                },
                {
                    view: gBufferTextures.color.createView(),
                    loadOp: loadOp,
                    clearValue: { r: 0.0, g: 0.0, b: 0.0, a: 1.0 },
                    storeOp: 'store',
                },
            ],
            depthStencilAttachment: {
                view: pipeline.depthView,
                depthClearValue: 1.0,
                depthLoadOp: loadOp,
                depthStoreOp: 'store',
            }
        }
        const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor)
        passEncoder.setPipeline(pipeline.pipeline)
        passEncoder.setBindGroup(0, pipeline.vsGroup)

        const { instanceBufferMap, instanceMap } = this;

        let renderedNum = 0;
        for (let i = 0, il = geometries.length; i < il; i++) {
            const geometry = geometries[i];
            const buffer = instanceBufferMap.get(geometry);
            if (buffer) {
                const num = instanceMap.get(geometry) ?? 0;
                if (num > 0) {
                    passEncoder.setVertexBuffer(0, buffer.vertex)
                    passEncoder.setIndexBuffer(buffer.index, 'uint16')
                    passEncoder.drawIndexed(geometry.indexCount, num, 0, 0, renderedNum)
                    renderedNum += num;
                }
            }
        }

        passEncoder.end()

    }

    getGbuffer() {
        return this.gbuffer;
    }

    public updateCamera() {
        const { device, instancePipelineMap, camera, lightPipeline } = this;
        const pps = instancePipelineMap.keys();

        let viewMatrix = camera.getViewMatrix();

        if (Array.isArray(viewMatrix)) {
            viewMatrix = new Float32Array(viewMatrix);
        }
        let projectionMatrix = camera.getProjectionMatrix();
        if (Array.isArray(projectionMatrix)) {
            projectionMatrix = new Float32Array(projectionMatrix);
        }
        pps.forEach((pipeline) => {
            device.queue.writeBuffer(pipeline.modelViewAllBuffer, 0, viewMatrix)
            device.queue.writeBuffer(pipeline.projectionBuffer, 0, projectionMatrix)
        });

        if (lightPipeline) {
            device.queue.writeBuffer(lightPipeline.modelViewAllBuffer, 0, viewMatrix)
        }
    }

    /**
     * 重新设置场景大小
     * @param size 场景大小
     */
    public resize(size: ISize) {
        this.size = size;
        const { device, instancePipelineMap, camera } = this;

        camera.resize(size);
        const gbuffer = createGBuffer(device, size);

        const pps = instancePipelineMap.keys();


        const depthView = gbuffer.depth.createView();
        pps.forEach((pipeline) => {
            pipeline.depthView = depthView;
            pipeline.gBufferTextures = gbuffer
            pipeline.depthTexture = gbuffer.depth
        });

        this.gbuffer = gbuffer;

        const oriGbufferArray = pps.toArray()
        if (oriGbufferArray.length > 0) {

            const oriGbuffer = oriGbufferArray[0].gBufferTextures;
            if (oriGbuffer) {
                const { position, normal, color, depth } = oriGbuffer;
                position.destroy();
                normal.destroy();
                color.destroy();
                depth.destroy();
            }
        }

        this.updateCamera();

        return gbuffer;
    }
    // 提供公共方法访问和修改内部状态
    public getDevice(): GPUDevice {
        return this.device;
    }

    public getContext(): GPUCanvasContext {
        return this.context;
    }

    public getFormat(): GPUTextureFormat {
        return this.format;
    }

    public getSize(): ISize {
        return this.size;
    }
}